"""
Copyright (R) @huawei.com, all rights reserved
-*- coding:utf-8 -*-
CREATED:  2020-6-04 20:12:13
MODIFIED: 2020-9-30 09:00:00
"""
# error code
ACL_ERROR_NONE = 0

# rule for mem
ACL_MEM_MALLOC_HUGE_FIRST = 0
ACL_MEM_MALLOC_HUGE_ONLY = 1
ACL_MEM_MALLOC_NORMAL_ONLY = 2

# rule for memory copy
ACL_MEMCPY_HOST_TO_HOST = 0
ACL_MEMCPY_HOST_TO_DEVICE = 1
ACL_MEMCPY_DEVICE_TO_HOST = 2
ACL_MEMCPY_DEVICE_TO_DEVICE = 3

# NUMPY data type
NPY_BOOL = 0
NPY_BYTE = 1
NPY_UBYTE = 2
NPY_SHORT = 3
NPY_USHORT = 4
NPY_INT = 5
NPY_UINT = 6
NPY_LONG = 7
NPY_ULONG = 8
NPY_LONGLONG = 9
NPY_ULONGLONG = 10

# data format
ACL_FORMAT_UNDEFINED = -1
ACL_FORMAT_NCHW = 0
ACL_FORMAT_NHWC = 1
ACL_FORMAT_ND = 2
ACL_FORMAT_NC1HWC0 = 3
ACL_FORMAT_FRACTAL_Z = 4

# data type
ACL_DT_UNDEFINED = -1
ACL_FLOAT = 0
ACL_FLOAT16 = 1
ACL_INT8 = 2
ACL_INT32 = 3
ACL_UINT8 = 4
ACL_INT16 = 6
ACL_UINT16 = 7
ACL_UINT32 = 8
ACL_INT64 = 9
ACL_UINT64 = 10
ACL_DOUBLE = 11
ACL_BOOL = 12

# dvpp pixel format
PIXEL_FORMAT_YUV_400 = 0  # YUV400 8bit
PIXEL_FORMAT_YUV_SEMIPLANAR_420 = 1  # YUV420SP NV12 8bit
PIXEL_FORMAT_YVU_SEMIPLANAR_420 = 2  # YUV420SP NV21 8bit
PIXEL_FORMAT_YUV_SEMIPLANAR_422 = 3  # YUV422SP NV12 8bit
PIXEL_FORMAT_YVU_SEMIPLANAR_422 = 4  # YUV422SP NV21 8bit
PIXEL_FORMAT_YUV_SEMIPLANAR_444 = 5  # YUV444SP NV12 8bit
PIXEL_FORMAT_YVU_SEMIPLANAR_444 = 6  # YUV444SP NV21 8bit
PIXEL_FORMAT_YUYV_PACKED_422 = 7  # YUV422P YUYV 8bit
PIXEL_FORMAT_UYVY_PACKED_422 = 8  # YUV422P UYVY 8bit
PIXEL_FORMAT_YVYU_PACKED_422 = 9  # YUV422P YVYU 8bit
PIXEL_FORMAT_VYUY_PACKED_422 = 10  # YUV422P VYUY 8bit
PIXEL_FORMAT_YUV_PACKED_444 = 11  # YUV444P 8bit
PIXEL_FORMAT_RGB_888 = 12  # RGB888
PIXEL_FORMAT_BGR_888 = 13  # BGR888
PIXEL_FORMAT_ARGB_8888 = 14  # ARGB8888
PIXEL_FORMAT_ABGR_8888 = 15  # ABGR8888
PIXEL_FORMAT_RGBA_8888 = 16  # RGBA8888
PIXEL_FORMAT_BGRA_8888 = 17  # BGRA8888
PIXEL_FORMAT_YUV_SEMI_PLANNER_420_10BIT = 18  # YUV420SP 10bit
PIXEL_FORMAT_YVU_SEMI_PLANNER_420_10BIT = 19  # YVU420sp 10bit
PIXEL_FORMAT_YVU_PLANAR_420 = 20  # YUV420P 8bit

# images format
IMG_EXT = ['.jpg', '.JPG', '.png', '.PNG', '.bmp', '.BMP', '.jpeg', '.JPEG']

import acl
import numpy

def check_ret(message, ret):
    if ret != ACL_ERROR_NONE:
        raise Exception("{} failed ret={}"
                        .format(message, ret))

class AscendContext(object):
    def __init__(self, device_id):
        self.device_id = device_id
        self.context = None

    def __enter__(self):
        ret = acl.init()
        check_ret("acl.init", ret)
        ret = acl.rt.set_device(self.device_id)
        check_ret("acl.rt.set_device", ret)
        context, ret = acl.rt.create_context(self.device_id)
        check_ret("acl.rt.create_context", ret)
        self.context = context
        return context

    def __exit__(self, exc_type, exc_val, exc_tb):
        ret = acl.rt.destroy_context(self.context)
        check_ret("acl.rt.destroy_context", ret)
        ret = acl.rt.reset_device(self.device_id)
        check_ret("acl.rt.reset_device", ret)
        ret = acl.finalize()
        check_ret("acl.finalize", ret)
        print('acl exit')


class NMat(object):
    def __init__(self, ptr_dev, data_size, width, height, stride_width, stride_height):
        self.ptr_dev = ptr_dev
        self.data_size = data_size
        self.width = width
        self.height = height
        self.stride_width = stride_width
        self.stride_height = stride_height

    def __del__(self):
        if self.ptr_dev != None:
            self.free()

    def free(self):
        ret = acl.media.dvpp_free(self.ptr_dev)
        check_ret("[NMat]acl.media.dvpp_free", ret)
        self.ptr_dev = None

    def numpy(self):
        host_ptr, ret = acl.rt.malloc_host(self.data_size)
        check_ret("[NMat]acl.rt.malloc_host", ret)
        ret = acl.rt.memcpy(host_ptr, self.data_size, self.ptr_dev, self.data_size, ACL_MEMCPY_DEVICE_TO_HOST)
        check_ret("[NMat]acl.rt.memcpy", ret)
        output = acl.util.ptr_to_numpy(host_ptr, (self.stride_height * 3 // 2, self.stride_width), ACL_INT8)
        return output

